from typing import Dict, List

import pandas as pd
from joblib import Parallel, delayed

from iflytech_assistant.client import instruct_with_profile

all_sheets = pd.read_excel(".vscode/胥纾.xls", sheet_name=None, header=None)


USER_PROMPT = """
请用你的语气表达
```
{user_input}
```
不要加引号，不要做多余的解释，请给出5条结果，以1开头，以5结束。如果是你，你会这样表达：

输出格式(xxx为直接改写结果,无需额外解释或标注):
1. xxx
2. xxx
3. xxx
4. xxx
5. xxx
"""


def run(df: pd.DataFrame):
    system_prompt = df.iloc[0, 1]
    start = 1

    outputs = [""]
    # 判断是否有示例
    line2 = df.iloc[1, 0]
    history: List[Dict[str, str]] = []
    if not pd.isna(line2) and "example" in line2:
        start = 2
        outputs.append("")

        ip, op = "", ""
        for l in df.iloc[1, 1].split("\n"):
            if l.startswith("输入"):
                if ip:
                    history.append(
                        {"role": "user", "content": USER_PROMPT.format(user_input=ip)}
                    )
                    ip = ""
                ip += l.replace("输入：", "")
            if l.startswith("输出"):
                if op:
                    history.append({"role": "assistant", "content": op})
                    op = ""
                op += l.replace("输出：", "")

        if ip:
            history.append(
                {"role": "user", "content": USER_PROMPT.format(user_input=ip)}
            )
        if op:
            history.append({"role": "assistant", "content": op})

    tasks = []
    for _, line in df.iloc[start:].iterrows():
        task = delayed(instruct_with_profile)(
            user_input=USER_PROMPT.format(user_input=line[1]),
            system_prompt=system_prompt,
            historys=history,
            temperature=0.0,
        )
        tasks.append(task)

    for response in Parallel(n_jobs=10, verbose=10)(tasks):
        outputs.append(response.llm_response)

    df["output"] = outputs


for sheet_name, df in all_sheets.items():
    print(f"Running sheet: {sheet_name}")
    try:
        run(df)
    except Exception as e:
        continue


# write to excel
with pd.ExcelWriter(".vscode/胥纾.xlsx") as writer:
    for sheet_name, df in all_sheets.items():
        df.to_excel(writer, sheet_name=sheet_name, index=False)
